#include "reg_alloc.h"
#include "cfg.h"
#include "reg.h"
using namespace reg_alloc;
using namespace ir;
using namespace live_var;
using namespace reg;
using namespace std;

static bool compare_end(LiveInterval* a, LiveInterval* b)
{ 
    return a->end < b->end; 
};

Function RegAllocator::Allocate(Function& func)
{
    cfg::CFG cfg = cfg::CFGGenerator().Generate(func);
    FuncLiveness func_liveness = LiveVarAnalyzer().Analyze(cfg);
    live_intervals = vector<LiveInterval>(func_liveness.var_index_map.size());
    active_intervals = 
        set<LiveInterval*, bool(*)(LiveInterval*, LiveInterval*)>(compare_end);
    used_regs = vector<bool>(32);
    call_sites.clear();
    call_sites_used_regs.clear();
    var_index_map = func_liveness.var_index_map;
    ComputeStartEnd(func, func_liveness);
    ComputeCallSite(func);
    ComputeCrossesCall();
    LinearScan();
    ComputeCallSiteUsedReg();
    Function new_func = func;
    ReplaceVar(new_func);
    AllocCallerSavedSlot(new_func);
    StoreLoadCallerSaved(new_func);
    AllocArgSlot(new_func);
    PassArg(new_func);
    return new_func;
}

void RegAllocator::ComputeStartEnd(Function& func, FuncLiveness& func_liveness)
{
    int instr_index = 0;
    for (Block& block : func.blocks)
    {
        BlockLiveness& block_liveness = 
            func_liveness.block_liveness_map[block.label];
        for (InstrLiveness& instr_liveness : block_liveness.instrs_liveness)
        {
            for (size_t i = 0; i < live_intervals.size(); ++i)
            {
                if (instr_liveness.def[i])
                { live_intervals[i].start = instr_index; }
                else if (instr_liveness.use[i])
                { live_intervals[i].end = instr_index; }
            }
            ++instr_index;
        }
    }
}

void RegAllocator::ComputeCallSite(Function& func)
{
    int instr_index = 0;
    for (Block& block: func.blocks)
    {
        for (Instruction& instr : block.instrs)
        {
            if (instr.op == OP_CALL) 
            { call_sites.push_back(instr_index); }
            ++instr_index;
        }
    }
}

void RegAllocator::ComputeCrossesCall()
{
    for (LiveInterval& live_interval : live_intervals)
    {
        for (int call_site : call_sites)
        {
            if (live_interval.start < call_site 
                && live_interval.end > call_site)
            {
                live_interval.crosses_call = true;
                break;
            }
        }
    }
}

void RegAllocator::LinearScan()
{
    int spill_index = 0;
    for (LiveInterval& live_interval : live_intervals)
    {
        ExpireInterval(live_interval.start);
        const vector<int>& regs_index = live_interval.crosses_call? 
            REG_CALLEE_SAVED_PRIOR : REG_TEMP_PRIOR;
        for (int reg_index : regs_index)
        {
            if (!used_regs[reg_index])
            {
                used_regs[reg_index] = true;
                live_interval.allocated = true;
                live_interval.reg_index = reg_index; 
                active_intervals.insert(&live_interval);
                break;
            }
        }
        if (!live_interval.allocated)
        {
            LiveInterval* active_interval = *active_intervals.rbegin();
            if (live_interval.end < active_interval->end)
            {
                live_interval.reg_index = active_interval->reg_index;
                live_interval.allocated = true;
                active_intervals.insert(&live_interval);
                active_interval->allocated = false;
                active_interval->spill_index = spill_index++;
                active_intervals.erase(active_interval);
            }
            else { live_interval.spill_index = spill_index++; }
        }
    }
}

void RegAllocator::ComputeCallSiteUsedReg()
{
    for (int call_site : call_sites)
    {
        vector<bool> call_site_used_regs(32);
        for (LiveInterval& live_interval : live_intervals)
        {
            if (live_interval.allocated 
                && live_interval.start < call_site 
                && live_interval.end > call_site)
            { call_site_used_regs[live_interval.reg_index] = true; }
        }
        call_sites_used_regs.push_back(move(call_site_used_regs));
    }
}

void RegAllocator::ReplaceVar(Function& func)
{
    for (Block& block : func.blocks)
    {
        vector<Instruction> new_instrs;
        for (Instruction& instr : block.instrs)
        {
            vector<Instruction> load_instrs;
            vector<Instruction> store_instrs;
            ReplaceInstrSrc(instr, load_instrs);
            ReplaceInstrDest(instr, store_instrs, func.alloc_instrs);
            new_instrs.insert(
                new_instrs.end(), load_instrs.begin(), load_instrs.end());
            new_instrs.push_back(instr);
            new_instrs.insert(
                new_instrs.end(), store_instrs.begin(), store_instrs.end());
        }
        block.instrs = move(new_instrs);
    }
}

void RegAllocator::AllocCallerSavedSlot(Function& func)
{
    int max_caller_saved_cnt = 0;
    for (vector<bool>& call_sites_used_reg : call_sites_used_regs)
    {
        int caller_saved_cnt = 0;
        for (int reg_index : REG_TEMP)
        {
            if (call_sites_used_reg[reg_index]) { ++caller_saved_cnt; }
        }
        if (max_caller_saved_cnt < caller_saved_cnt) 
        { max_caller_saved_cnt = caller_saved_cnt; }
    }
    vector<Instruction> alloc_instrs;
    for (int i = 0; i < max_caller_saved_cnt; ++i)
    {
        Instruction alloc_instr;
        alloc_instr.op = OP_ALLOC;
        alloc_instr.dest = "%caller_saved_" + to_string(i);
        alloc_instrs.push_back(move(alloc_instr));
    }
    func.alloc_instrs.insert(
        func.alloc_instrs.begin(), alloc_instrs.begin(), alloc_instrs.end());
}

void RegAllocator::StoreLoadCallerSaved(Function& func)
{
    int call_site_index = 0;
    for (Block& block: func.blocks)
    {
        vector<Instruction> new_instrs;
        for (Instruction& instr : block.instrs)
        {
            vector<Instruction> store_instrs;
            vector<Instruction> load_instrs;
            if (instr.op == OP_CALL)
            {
                if (!call_sites_used_regs.empty())
                {
                    StoreLoadInstrCallerSaved(
                        call_sites_used_regs[call_site_index], 
                        store_instrs, load_instrs);
                    ++call_site_index;
                }
            }
            new_instrs.insert(
                new_instrs.end(), store_instrs.begin(), store_instrs.end());
            new_instrs.push_back(instr);
            new_instrs.insert(
                new_instrs.end(), load_instrs.begin(), load_instrs.end());
        }
        block.instrs = move(new_instrs);
    }
}

void RegAllocator::AllocArgSlot(Function& func)
{
    size_t max_arg_cnt = 0;
    for (Block& block : func.blocks)
    {
        for (Instruction& instr : block.instrs)
        {
            if (instr.op == OP_CALL)
            { 
                if (max_arg_cnt < instr.srcs.size() - 1) 
                { max_arg_cnt = instr.srcs.size() - 1; }
            }
        }
    }
    for (int i = max_arg_cnt - 1; i >= (int)REG_ARG.size(); --i)
    {
        Instruction alloc_instr;
        alloc_instr.op = OP_ALLOC;
        alloc_instr.dest = "%arg_" + to_string(i);
        func.alloc_instrs.push_back(move(alloc_instr));
    }
}

void RegAllocator::PassArg(Function& func)
{
    for (Block& block : func.blocks)
    {
        vector<Instruction> new_instrs;
        for (Instruction& instr : block.instrs)
        {
            if (instr.op == OP_CALL)
            {
                PassInstrExtraArg(instr, new_instrs);
                PassInstrArg(instr, new_instrs);
                instr.srcs.resize(1);
            }
            new_instrs.push_back(instr);
        }
        block.instrs = move(new_instrs);
    }
}

void RegAllocator::ExpireInterval(int instr_index)
{
    for (auto it = active_intervals.begin(); it != active_intervals.end();)
    {
        LiveInterval* active_interval = *it;
        if (active_interval->end <= instr_index)
        {
            used_regs[active_interval->reg_index] = false;
            it = active_intervals.erase(it);
        }
        else { break; }
    }
}

void RegAllocator::ReplaceInstrSrc(
    Instruction& instr, vector<Instruction>& load_instrs)
{
    if (instr.op != OP_CALL)
    {
        auto spill_reg_index = REG_ARG.begin();
        for (size_t i = 0; i < instr.srcs.size(); ++i)
        {
            if (holds_alternative<string>(instr.srcs[i]))
            {
                string var_name = get<string>(instr.srcs[i]);
                if (var_index_map.count(var_name))
                {
                    int var_index = var_index_map[var_name];
                    LiveInterval& live_interval = live_intervals[var_index];
                    if (live_interval.allocated)
                    { 
                        instr.srcs[i] = 
                            "x" + to_string(live_interval.reg_index); 
                    }
                    else
                    {
                        string reg_name = 
                            "x" + to_string(*(spill_reg_index++));
                        instr.srcs[i] = reg_name;
                        Instruction load_instr;
                        load_instr.op = OP_LOAD;
                        load_instr.dest = reg_name;
                        string spill_slot = 
                            "%spill_" + to_string(live_interval.spill_index);
                        load_instr.srcs.push_back(move(spill_slot));
                        load_instrs.push_back(move(load_instr));
                    }
                }
            }
        }
    }
}

void RegAllocator::ReplaceInstrDest(
    Instruction& instr, vector<Instruction>& store_instrs, 
    vector<Instruction>& alloc_instrs)
{
    if (instr.dest != "")
    {
        int var_index = var_index_map[instr.dest];
        LiveInterval& live_interval = live_intervals[var_index];
        if (live_interval.allocated)
        { instr.dest = "x" + to_string(live_interval.reg_index); }
        else
        {
            string reg_name = "x" + to_string(REG_ARG.front());
            instr.dest = reg_name;
            Instruction store_instr;
            store_instr.op = OP_STORE;
            store_instr.dest = "";
            store_instr.srcs.push_back(move(reg_name));
            string spill_slot = 
                "%spill_" + to_string(live_interval.spill_index);
            store_instr.srcs.push_back(spill_slot);
            store_instrs.push_back(move(store_instr));
            Instruction alloc_instr;
            alloc_instr.op = OP_ALLOC;
            alloc_instr.dest = spill_slot;
            alloc_instrs.push_back(move(alloc_instr));
        }
    }
}

void RegAllocator::StoreLoadInstrCallerSaved(
    vector<bool>& call_site_used_regs, vector<Instruction>& store_instrs, 
    vector<Instruction>& load_instrs)
{
    int caller_saved_index = 0;
    for (int reg_index : REG_TEMP)
    {
        if (call_site_used_regs[reg_index])
        {
            Instruction store_instr;
            store_instr.op = OP_STORE;
            store_instr.dest = "";
            string reg_name = "x" + to_string(reg_index);
            store_instr.srcs.push_back(reg_name);
            string caller_saved_slot = 
                "%caller_saved_" + to_string(caller_saved_index++);
            store_instr.srcs.push_back(caller_saved_slot);
            store_instrs.push_back(move(store_instr));
            Instruction load_instr;
            load_instr.op = OP_LOAD;
            load_instr.dest = reg_name;
            load_instr.srcs.push_back(caller_saved_slot);
            load_instrs.push_back(move(load_instr));                       
        }
    }
}

void RegAllocator::PassInstrArg(
    Instruction& instr, vector<Instruction>& new_instrs)
{
    for (size_t i = 0; i < REG_ARG.size() && i + 1 < instr.srcs.size(); ++i)
    {
        Instruction new_instr;
        new_instr.dest = "x" + to_string(REG_ARG[i]); 
        if (holds_alternative<string>(instr.srcs[i+1]))
        {
            string var_name = get<string>(instr.srcs[i+1]);
            int var_index = var_index_map[var_name];
            LiveInterval& live_interval = live_intervals[var_index];
            if (live_interval.allocated)
            {
                new_instr.op = OP_MV;
                string reg_name = "x" + to_string(live_interval.reg_index);
                new_instr.srcs.push_back(move(reg_name));           
            }
            else
            {
                new_instr.op = OP_LOAD;
                string spill_slot = 
                    "%spill_" + to_string(live_interval.spill_index);
                new_instr.srcs.push_back(move(spill_slot));
            }
        }
        else
        {
            new_instr.op = OP_MV;
            new_instr.srcs.push_back(instr.srcs[i+1]);
        }
        new_instrs.push_back(move(new_instr));
    }
}

void RegAllocator::PassInstrExtraArg(
    Instruction& instr, vector<Instruction>& new_instrs)
{
    for (size_t i = REG_ARG.size(); i + 1 < instr.srcs.size(); ++i)
    {
        Instruction store_instr;
        store_instr.op = OP_STORE;
        store_instr.dest = "";
        string arg_slot = "%arg_" + to_string(i);
        if (holds_alternative<string>(instr.srcs[i+1]))
        {
            string var_name = get<string>(instr.srcs[i+1]);
            int var_index = var_index_map[var_name];
            LiveInterval& live_interval = live_intervals[var_index];
            if (live_interval.allocated)
            {
                string reg_name = "x" + to_string(live_interval.reg_index);
                store_instr.srcs.push_back(move(reg_name));
                store_instr.srcs.push_back(move(arg_slot));             
            }
            else
            {
                Instruction load_instr;
                load_instr.op = OP_LOAD;
                string reg_name = "x" + to_string(REG_TEMP.front());
                load_instr.dest = reg_name;
                string spill_slot = 
                    "%spill_" + to_string(live_interval.spill_index);
                load_instr.srcs.push_back(move(spill_slot));
                new_instrs.push_back(move(load_instr));
                store_instr.srcs.push_back(move(reg_name));
                store_instr.srcs.push_back(move(arg_slot));       
            }
        }
        else
        {
            store_instr.srcs.push_back(instr.srcs[i+1]);
            store_instr.srcs.push_back(move(arg_slot));
        }
        new_instrs.push_back(move(store_instr));
    }
}